import numpy as np
import gym
from gym import spaces
from envs.env_core import EnvCore


class UAV_CPP_Utils:

    def __init__(self):
        pass

    @staticmethod
    def _find_nearest_uncovered(self, agent_id):
        """
        寻找距离当前无人机最近的未覆盖区域
        """
        x, y = self.uav_positions[agent_id]
        min_dist = float('inf')
        target = None
        
        # 遍历地图寻找未覆盖区域
        for i in range(self.grid_size):
            for j in range(self.grid_size):
                if self.uncovered_map[i, j] == 1:  # 未覆盖区域
                    dist = abs(x - i) + abs(y - j)  # 曼哈顿距离
                    if dist < min_dist:
                        min_dist = dist
                        target = (i, j)
        
        return target

    @staticmethod
    def _plan_path_to_target(self, agent_id, target_pos):
        """
        使用A*算法为无人机规划从当前位置到目标位置的路径
        添加防震荡机制：目标到达容忍度和路径平滑处理
        
        Args:
            agent_id: 无人机ID
            target_pos: 目标位置 (x, y)
            
        Returns:
            next_action: 下一步应该执行的动作 (0-上, 1-右, 2-下, 3-左)
        """
        if target_pos is None:
            return None
        
        # 获取当前位置
        start_pos = tuple(self.uav_positions[agent_id])
        target_pos = tuple(target_pos)
        
        # 计算到目标的距离
        dist_to_target = abs(start_pos[0] - target_pos[0]) + abs(start_pos[1] - target_pos[1])
        
        # 防震荡：增加目标到达的容忍度
        # 如果距离目标只有1格，且目标位置没有障碍物，直接返回朝向目标的动作
        if dist_to_target == 1:
            dx = target_pos[0] - start_pos[0]
            dy = target_pos[1] - start_pos[1]
            if dx == 1:
                return 2  # 下
            elif dx == -1:
                return 0  # 上
            elif dy == 1:
                return 1  # 右
            elif dy == -1:
                return 3  # 左
        
        # 如果已经到达目标，返回None
        if start_pos == target_pos:
            return None
        
        # 获取无人机历史路径，用于防止来回震荡
        recent_positions = []
        if hasattr(self, 'path_history') and len(self.path_history) > agent_id:
            # 获取最近的几个位置
            path = self.path_history[agent_id]
            if len(path) >= 3:
                recent_positions = path[-3:]
        
        # A*算法
        # 定义启发式函数 - 曼哈顿距离
        def heuristic(a, b):
            return abs(a[0] - b[0]) + abs(a[1] - b[1])
        
        # 定义可能的移动方向 (上, 右, 下, 左)
        directions = [(-1, 0), (0, 1), (1, 0), (0, -1)]
        direction_to_action = {(-1, 0): 0, (0, 1): 1, (1, 0): 2, (0, -1): 3}
        
        # 初始化开放列表和关闭列表
        open_set = []
        closed_set = set()
        
        # 使用优先队列实现开放列表
        import heapq
        
        # 记录从起点到每个节点的最小成本和父节点
        g_score = {start_pos: 0}
        f_score = {start_pos: heuristic(start_pos, target_pos)}
        came_from = {}
        
        # 将起点加入开放列表
        heapq.heappush(open_set, (f_score[start_pos], start_pos))
        
        while open_set:
            # 获取f值最小的节点
            _, current = heapq.heappop(open_set)
            
            # 如果到达目标，构建路径并返回下一步动作
            if current == target_pos:
                # 回溯路径
                path = [current]
                while current in came_from:
                    current = came_from[current]
                    path.append(current)
                path.reverse()
                
                # 确定下一步动作
                if len(path) > 1:
                    next_pos = path[1]
                    
                    # 防震荡：检查是否形成了来回震荡的模式
                    if len(recent_positions) >= 2:
                        # 检查是否在两个位置之间来回移动
                        if next_pos == recent_positions[-2] and start_pos == recent_positions[-1]:
                            # 检测到震荡，尝试选择另一个方向
                            alternative_directions = []
                            for dx, dy in directions:
                                alt_pos = (start_pos[0] + dx, start_pos[1] + dy)
                                # 确保替代位置在地图内，不是障碍物，且不是刚才的位置
                                if (0 <= alt_pos[0] < self.grid_size and 
                                    0 <= alt_pos[1] < self.grid_size and 
                                    self.obstacle_grid[alt_pos[0], alt_pos[1]] == 0 and 
                                    alt_pos != recent_positions[-2]):
                                    # 计算到目标的距离
                                    alt_dist = abs(alt_pos[0] - target_pos[0]) + abs(alt_pos[1] - target_pos[1])
                                    alternative_directions.append((alt_dist, alt_pos))
                            
                            # 如果有替代方向，选择距离目标最近的一个
                            if alternative_directions:
                                alternative_directions.sort()  # 按距离排序
                                next_pos = alternative_directions[0][1]  # 选择最近的
                    
                    dx = next_pos[0] - start_pos[0]
                    dy = next_pos[1] - start_pos[1]
                    return direction_to_action[(dx, dy)]
                return None
            
            # 将当前节点加入关闭列表
            closed_set.add(current)
            
            # 检查所有相邻节点
            for dx, dy in directions:
                neighbor = (current[0] + dx, current[1] + dy)
                
                # 检查边界
                if not (0 <= neighbor[0] < self.grid_size and 0 <= neighbor[1] < self.grid_size):
                    continue
                
                # 检查障碍物
                if self.obstacle_grid[neighbor[0], neighbor[1]] == 1:
                    continue
                
                # 如果邻居在关闭列表中，跳过
                if neighbor in closed_set:
                    continue
                
                # 计算从起点经过当前节点到邻居的成本
                tentative_g_score = g_score[current] + 1
                
                # 防震荡：如果这个位置是最近访问过的，增加成本以避免来回移动
                if neighbor in recent_positions:
                    tentative_g_score += 2.0
                
                # 如果邻居不在开放列表中，或者找到了更好的路径
                if neighbor not in g_score or tentative_g_score < g_score[neighbor]:
                    # 更新邻居的信息
                    came_from[neighbor] = current
                    g_score[neighbor] = tentative_g_score
                    f_score[neighbor] = tentative_g_score + heuristic(neighbor, target_pos)
                    
                    # 将邻居加入开放列表
                    heapq.heappush(open_set, (f_score[neighbor], neighbor))
        
        # 如果没有找到路径，使用简单的贪心策略，但避免来回震荡
        x, y = start_pos
        target_x, target_y = target_pos
        
        # 优先考虑x方向还是y方向的移动
        # 这里添加一些随机性，避免在相似情况下总是做出相同决策
        import random
        if abs(x - target_x) > abs(y - target_y) or (abs(x - target_x) == abs(y - target_y) and random.random() < 0.5):
            # 优先考虑x方向
            if x < target_x:
                return 2  # 下
            elif x > target_x:
                return 0  # 上
            elif y < target_y:
                return 1  # 右
            elif y > target_y:
                return 3  # 左
        else:
            # 优先考虑y方向
            if y < target_y:
                return 1  # 右
            elif y > target_y:
                return 3  # 左
            elif x < target_x:
                return 2  # 下
            elif x > target_x:
                return 0  # 上
        
        return None

    @staticmethod
    def _find_best_uncovered_target(self, agent_id):
        """
        寻找最佳的未覆盖目标，考虑其他智能体的目标
        添加防震荡机制：目标稳定性和历史信息
        """
        # 如果是主从模式，且当前智能体是主机，则分配任务给从机
        if self.leader_follower_mode and agent_id == self.leader_id:
            self._leader_assign_tasks()
        
        # 如果是从机，且主机已经分配了目标，则使用主机分配的目标
        if self.leader_follower_mode and agent_id != self.leader_id:
            # 检查主机是否给当前从机分配了目标
            leader_command_target = self.communication_data[self.leader_id].get('command_target', None)
            if leader_command_target is not None and isinstance(leader_command_target, dict) and str(agent_id) in leader_command_target:
                assigned_target = leader_command_target[str(agent_id)]
                if assigned_target is not None:
                    # 存储目标位置，供通信和可视化使用
                    self.communication_data[agent_id]['target'] = assigned_target
                    
                    # 计算到达目标的下一步动作
                    next_action = self._plan_path_to_target(agent_id, assigned_target)
                    self.communication_data[agent_id]['next_action'] = next_action
                    
                    return assigned_target
        
        # 获取当前位置和当前目标
        current_pos = tuple(self.uav_positions[agent_id])
        current_target = self.communication_data[agent_id].get('target', None)
        
        # 确保current_target是元组类型，如果是numpy数组则转换
        if current_target is not None:
            current_target = tuple(current_target)
        
        # 防震荡：如果已有目标且距离较近，保持当前目标
        # 这样可以避免在接近目标时频繁切换目标
        if current_target is not None:
            dist_to_current_target = abs(current_pos[0] - current_target[0]) + abs(current_pos[1] - current_target[1])
            # 如果距离当前目标很近（小于3格），且目标仍未覆盖，继续使用当前目标
            if dist_to_current_target < 3 and self.uncovered_map[current_target[0], current_target[1]] == 1:
                return current_target
        
        # 获取其他智能体的目标
        other_targets = []
        for i in range(self.agent_num):
            if i != agent_id and self.communication_data[i]['target'] is not None:
                other_targets.append(self.communication_data[i]['target'])
        
        # 寻找未覆盖区域
        uncovered_positions = []
        for x in range(self.grid_size):
            for y in range(self.grid_size):
                if self.uncovered_map[x, y] == 1:  # 未覆盖且非障碍物
                    uncovered_positions.append((x, y))
        
        if not uncovered_positions:
            return None
        
        # 计算每个未覆盖位置的得分
        # 得分 = -距离 - 与其他智能体目标的接近度 + 目标稳定性奖励
        best_score = float('-inf')
        best_target = None
        
        for pos in uncovered_positions:
            # 计算距离分数 (负距离，越近越好)
            distance = abs(pos[0] - current_pos[0]) + abs(pos[1] - current_pos[1])
            distance_score = -distance
            
            # 计算与其他智能体目标的接近度 (负接近度，越远越好)
            proximity_score = 0
            for target in other_targets:
                if target is not None:
                    target_distance = abs(pos[0] - target[0]) + abs(pos[1] - target[1])
                    # 如果非常接近其他智能体的目标，给予较大的负分
                    if target_distance < 3:  # 距离小于3个格子
                        proximity_score -= (3 - target_distance) * 5  # 越近惩罚越大
            
            # 计算电量因素 (电量低时倾向于选择更近的目标)
            battery_factor = 1.0
            if self.uav_battery[agent_id] < self.low_battery_threshold:
                battery_factor = self.uav_battery[agent_id] / self.low_battery_threshold
            
            # 目标稳定性奖励 - 如果是当前目标，给予额外奖励以增加稳定性
            stability_score = 0
            if current_target is not None and pos == current_target:  # 这里比较元组
                stability_score = 5.0  # 较大的稳定性奖励，减少目标切换
            
            # 综合得分，电量低时更重视距离
            total_score = distance_score * battery_factor + proximity_score + stability_score
            
            if total_score > best_score:
                best_score = total_score
                best_target = pos
        
        # 添加：如果找到目标，计算下一步动作
        if best_target is not None:
            # 存储目标位置，供通信和可视化使用
            self.communication_data[agent_id]['target'] = best_target
            
            # 计算到达目标的下一步动作
            next_action = self._plan_path_to_target(agent_id, best_target)
            self.communication_data[agent_id]['next_action'] = next_action
        
        return best_target
    
    @staticmethod
    def _leader_assign_tasks(self):
        """
        主机分配任务给从机 - 改进版
        考虑工作负载平衡、电量状态和区域划分
        """
        # 只有主机才能分配任务
        if not self.leader_follower_mode or self.leader_id is None:
            return
        
        # 获取所有未覆盖区域
        uncovered_positions = []
        for x in range(self.grid_size):
            for y in range(self.grid_size):
                if self.uncovered_map[x, y] == 1:  # 未覆盖且非障碍物
                    uncovered_positions.append((x, y))
        
        if not uncovered_positions:
            return
        
        # 获取所有从机的位置和状态
        followers = []
        for i in range(self.agent_num):
            if i != self.leader_id:  # 不是主机
                # 考虑所有从机，但记录其电量状态
                followers.append({
                    'id': i,
                    'position': tuple(self.uav_positions[i]),
                    'battery': self.uav_battery[i],
                    # 计算当前工作负载 - 已覆盖区域数量
                    'workload': np.sum(self.agent_coverage_maps[i]),
                    # 电量状态分类
                    'battery_status': 'critical' if self.uav_battery[i] <= self.critical_battery_threshold else
                                     'low' if self.uav_battery[i] <= self.low_battery_threshold else 'normal'
                })
        
        # 如果没有可用的从机，则返回
        if not followers:
            return
        
        # 初始化分配结果
        assignments = {str(follower['id']): None for follower in followers}
        
        # 特殊处理：如果有从机电量处于危险状态，优先为其分配最近的充电站
        for follower in followers:
            if follower['battery_status'] == 'critical':
                # 寻找最近的充电站
                follower_pos = follower['position']
                min_dist = float('inf')
                closest_station = None
                
                for station in self.charging_stations:
                    dist = abs(station[0] - follower_pos[0]) + abs(station[1] - follower_pos[1])
                    if dist < min_dist:
                        min_dist = dist
                        closest_station = tuple(station)
                
                if closest_station is not None:
                    assignments[str(follower['id'])] = closest_station
        
        # 对剩余未分配的从机进行任务分配
        # 首先，将地图划分为区域
        if len(uncovered_positions) >= len(followers):
            # 使用K-means聚类算法将未覆盖区域分组
            # 简化版K-means，适用于小规模问题
            
            # 过滤掉已经分配了目标的从机
            unassigned_followers = [f for f in followers if assignments[str(f['id'])] is None]
            if not unassigned_followers:
                # 所有从机都已分配任务
                self.communication_data[self.leader_id]['command_target'] = assignments
                return
                
            # 聚类数量 = 未分配从机数量
            k = len(unassigned_followers)
            
            # 初始化聚类中心 - 使用最远点启发式
            centers = []
            if uncovered_positions:
                # 选择第一个中心点 - 随机选择
                first_center = uncovered_positions[np.random.randint(0, len(uncovered_positions))]
                centers.append(first_center)
                
                # 选择剩余的中心点 - 每次选择距离已有中心点最远的点
                for _ in range(1, k):
                    if len(centers) >= len(uncovered_positions):
                        break
                        
                    max_min_dist = -1
                    next_center = None
                    
                    for pos in uncovered_positions:
                        if pos in centers:
                            continue
                            
                        # 计算到所有已有中心的最小距离
                        min_dist = float('inf')
                        for center in centers:
                            dist = abs(pos[0] - center[0]) + abs(pos[1] - center[1])
                            min_dist = min(min_dist, dist)
                        
                        # 更新最大最小距离
                        if min_dist > max_min_dist:
                            max_min_dist = min_dist
                            next_center = pos
                    
                    if next_center:
                        centers.append(next_center)
            
            # 如果中心点不足，使用从机位置作为补充
            while len(centers) < k:
                for follower in unassigned_followers:
                    if len(centers) >= k:
                        break
                    if follower['position'] not in centers:
                        centers.append(follower['position'])
            
            # 执行聚类
            clusters = [[] for _ in range(k)]
            
            # 将每个未覆盖位置分配到最近的聚类中心
            for pos in uncovered_positions:
                min_dist = float('inf')
                closest_cluster = 0
                
                for i, center in enumerate(centers):
                    dist = abs(pos[0] - center[0]) + abs(pos[1] - center[1])
                    if dist < min_dist:
                        min_dist = dist
                        closest_cluster = i
                
                clusters[closest_cluster].append(pos)
            
            # 根据从机状态为每个从机分配最合适的聚类
            # 1. 按电量状态和工作负载排序从机
            # 电量低的从机优先选择较小的聚类，工作负载大的从机优先选择较小的聚类
            
            # 计算每个聚类的大小和到各个从机的平均距离
            cluster_metrics = []
            for i, cluster in enumerate(clusters):
                if not cluster:  # 空聚类
                    cluster_metrics.append({
                        'id': i,
                        'size': 0,
                        'distances': {str(f['id']): float('inf') for f in unassigned_followers}
                    })
                    continue
                    
                # 计算聚类中心
                center_x = sum(p[0] for p in cluster) / len(cluster)
                center_y = sum(p[1] for p in cluster) / len(cluster)
                cluster_center = (center_x, center_y)
                
                # 计算到各个从机的距离
                distances = {}
                for follower in unassigned_followers:
                    follower_pos = follower['position']
                    dist = abs(follower_pos[0] - cluster_center[0]) + abs(follower_pos[1] - cluster_center[1])
                    distances[str(follower['id'])] = dist
                
                cluster_metrics.append({
                    'id': i,
                    'size': len(cluster),
                    'center': cluster_center,
                    'distances': distances
                })
            
            # 按大小排序聚类（从小到大）
            cluster_metrics.sort(key=lambda x: x['size'])
            
            # 为每个从机分配聚类 - 使用匈牙利算法的简化版本
            # 根据从机的电量状态、工作负载和到聚类的距离计算成本矩阵
            cost_matrix = []
            for follower in unassigned_followers:
                follower_costs = []
                for cluster in cluster_metrics:
                    # 基础成本 = 距离
                    base_cost = cluster['distances'][str(follower['id'])]
                    
                    # 电量因子 - 电量越低，距离因子越重要
                    battery_factor = 1.0
                    if follower['battery_status'] == 'low':
                        battery_factor = 2.0  # 低电量时距离权重加倍
                    elif follower['battery_status'] == 'critical':
                        battery_factor = 4.0  # 危险电量时距离权重更大
                    
                    # 工作负载因子 - 已完成工作量越大，分配的聚类应越小
                    workload_factor = follower['workload'] / self.total_cells
                    
                    # 聚类大小因子 - 聚类越大，成本越高
                    size_factor = cluster['size'] / max(1, len(uncovered_positions))
                    
                    # 综合成本
                    total_cost = base_cost * battery_factor + 100 * workload_factor * size_factor
                    follower_costs.append(total_cost)
                
                cost_matrix.append(follower_costs)
            
            # 使用贪心算法分配聚类
            assigned_clusters = set()
            for i, follower in enumerate(unassigned_followers):
                # 找到成本最低的未分配聚类
                min_cost = float('inf')
                best_cluster = None
                
                for j, cost in enumerate(cost_matrix[i]):
                    if j not in assigned_clusters and cost < min_cost:
                        min_cost = cost
                        best_cluster = j
                
                if best_cluster is not None:
                    # 分配聚类
                    assigned_clusters.add(best_cluster)
                    
                    # 从聚类中选择最佳目标点
                    cluster = clusters[cluster_metrics[best_cluster]['id']]
                    if cluster:
                        # 选择距离从机最近的点
                        follower_pos = follower['position']
                        min_dist = float('inf')
                        best_target = None
                        
                        for pos in cluster:
                            dist = abs(pos[0] - follower_pos[0]) + abs(pos[1] - follower_pos[1])
                            if dist < min_dist:
                                min_dist = dist
                                best_target = pos
                        
                        if best_target:
                            assignments[str(follower['id'])] = best_target
        else:
            # 如果未覆盖区域少于从机数量，使用简单的贪心分配
            # 按电量从高到低排序从机
            unassigned_followers = [f for f in followers if assignments[str(f['id'])] is None]
            unassigned_followers.sort(key=lambda x: x['battery'], reverse=True)
            
            # 已分配的位置
            assigned_positions = set()
            
            for follower in unassigned_followers:
                follower_id = follower['id']
                follower_pos = follower['position']
                
                # 找到距离当前从机最近且未分配的未覆盖区域
                best_pos = None
                best_dist = float('inf')
                
                for pos in uncovered_positions:
                    if pos not in assigned_positions:
                        dist = abs(pos[0] - follower_pos[0]) + abs(pos[1] - follower_pos[1])
                        
                        # 考虑电量因素
                        if follower['battery_status'] == 'low':
                            # 低电量时更看重距离
                            dist *= 1.5
                        elif follower['battery_status'] == 'critical':
                            # 危险电量时极度看重距离
                            dist *= 3.0
                            
                        if dist < best_dist:
                            best_dist = dist
                            best_pos = pos
                
                if best_pos is not None:
                    assignments[str(follower_id)] = best_pos
                    assigned_positions.add(best_pos)
        
        # 更新主机的命令目标
        self.communication_data[self.leader_id]['command_target'] = assignments

    @staticmethod
    def _calculate_observation_dim(self, agent_num=None):
        """
        计算观测空间维度
        """
        if agent_num is None:
            agent_num = self.agent_num
            
        # 2(位置) + 1(电量) + 8(周围栅格) + 2(充电站位置) + 1(覆盖率) + 1(是否在充电站) 
        # + 1(障碍物指示器) + 1(电量状态) + 1(充电站距离) + 1(其他智能体覆盖比例) 
        # + 1(未覆盖区域比例) + 1(自己覆盖区域比例) + 8(周围未覆盖详情) + 8(周围自己覆盖详情)
        # + 1(归一化步数) + 1(角色标识) + 9*(agent_num-1)(通信信息)
        base_dim = 2 + 1 + 8 + 2 + 1 + 1 + 1 + 1 + 1 + 1 + 1 + 1 + 8 + 8 + 1 + 1  # 基础维度(39)
        comm_dim = 9 * (agent_num - 1)  # 通信维度 - 增加了角色和命令信息
        return base_dim + comm_dim

    @staticmethod
    def _generate_charging_stations(self, required_stations):
        """
        根据需要的充电站数量生成充电站位置
        
        Args:
            required_stations: 需要的充电站数量
        """
        # 创建新的充电站数组
        self.charging_stations = np.zeros((required_stations, 2), dtype=int)
        
        # 生成随机充电站位置，确保不在障碍物上
        used_positions = set()
        
        for i in range(required_stations):
            while True:
                # 随机生成位置
                new_pos = [np.random.randint(0, self.grid_size), np.random.randint(0, self.grid_size)]
                
                
                # 检查位置是否已被使用且不在障碍物上
                if (tuple(new_pos) not in used_positions and 
                    self.obstacle_grid[new_pos[0], new_pos[1]] == 0):
                    self.charging_stations[i] = np.array(new_pos)
                    used_positions.add(tuple(new_pos))
                    break
        
      #   print(f"充电站位置已更新: {self.charging_stations}")
